package TestPackage;

import DataSetPackage.DataSet;
import DataSetPackage.DataSetStorage;
import DataSetPackage.ProductView;
import DataSetPackage.SimpleView;
import DataSetPackage.View;
import EstimatorPackage.EmpSuffix;
import EstimatorPackage.EstimatorExp;
import EstimatorPackage.IndependentProductEst;
import EstimatorPackage.LaPlaceSuffix;
import EstimatorPackage.SimSuffixEstimatorExp;
import HierarchyPackage.Hierarchy;
import HierarchyPackage.ProductHierarchy.ProductHierarchy;
import HierarchyPackage.AffixHierarchies.SuffixHierarchy;
import HierarchyPackage.HClass;
import HierarchyPackage.ProductHierarchy.ProductHClass;
import HierarchyPackage.SimpleHierarchy.SimpleHierarchy;
import HierarchyPackage.SimpleHierarchy.SimpleHierarchyMem;
import OptionsManager.DiscountOptions;
import OptionsManager.TrainOptions;
import OrderedCoverPackage.BaseHPM;
import OrderedCoverPackage.NSuffixOCEstimator;
import OrderedCoverPackage.OrderedCoverEstimator;
import java.io.BufferedReader;
import java.io.FileReader;
import java.io.FileWriter;
import java.io.PrintWriter;
import java.util.ArrayList;
import java.util.List;
import org.osdtsystem.utils.MemoryStorage;
import org.osdtsystem.utils.Storage;
import util.IO;
import util.WordPOS;

/**
 *
 * @author Martin Haulrich
 */
public class Main {

    static String newline = System.getProperty("line.separator");

    public static final String SOSPOS = "<SOS>";

    public static final String EOSPOS = "<EOS>";

    /** Constructor */
    public Main() {
    }

    /**
     * @param args the command line arguments
     */
    public static void main(String[] args) throws Exception {

        if (false) {
            SimpleHierarchy POSHier = new SimpleHierarchyMem("output/SUC.hier");
            System.err.println(POSHier);

            Storage storage = new MemoryStorage();
            DataSet ds = new DataSetStorage(storage);

            List<List<WordPOS>> train = IO.readSUC("/Users/mwh/Java/HMM-Tagger/train.txt");

            List<String> priorTrain = new ArrayList<String>();

            for (List<WordPOS> sen : train) {
                List<String> dp = new ArrayList<String>(2);
                dp.add(SOSPOS);
                dp.add(sen.get(0).POS);
                ds.add(dp);
                for (int i = 1; i < sen.size(); i++) {
                    dp = new ArrayList<String>(2);
                    dp.add(sen.get(i - 1).POS);
                    dp.add(sen.get(i).POS);
                    ds.add(dp);
                }
                dp = new ArrayList<String>(2);
                dp.add(sen.get(sen.size() - 1).POS);
                dp.add(EOSPOS);
                ds.add(dp);

                priorTrain.add(SOSPOS);
                for (int i = 0; i < sen.size(); i++) {
                    priorTrain.add(sen.get(i).POS);
                }
                priorTrain.add(EOSPOS);
            }

            EmpPOS es = new EmpPOS(priorTrain);

            EmpPOS ls = new LaPlacePOS(es);
            // Create estimator to handle expression
            EstimatorExp ees = new SimPOSEstimatorExp(es, POSHier, 100);
            EstimatorExp els = new SimPOSEstimatorExp(ls, POSHier, 100);

            EstimatorExp[] ee = {els, ees};

            // Create product-estimator
            IndependentProductEst ie = new IndependentProductEst(ee);

            EstimatorExp pe = ie;

            System.err.println("DatasetSize: " + ds.size());

            // Data for hierarchy is at index 0.

            SimpleView posdv1 = new SimpleView(0, ds);
            SimpleView posdv2 = new SimpleView(1, ds);

            // Create productView
//            View[] views = {sufdv1, sufdv2};
            View[] views = {posdv1, posdv2};
            ProductView pv = new ProductView(ds, views);


            System.err.println(ds.getFeatures(4));
            System.err.println(es.getEstimate("PP"));
            System.err.println(ees.getEstimate(4, posdv1));
            System.err.println(ees.getEstimate(4, posdv2));
            System.err.println(ie.getEstimate(4, pv));


            for (Integer did : ds.getIDs()) {
//                System.err.println(ds.getFeatures(did) + "\t" + pe.getEstimate(did, pv));
            }

//            Hierarchy[] hiers = {sufhier, sufhier};
            Hierarchy[] hiers = {POSHier, POSHier};

            // And now we make the suf x suf producthierarchy
            ProductHierarchy phier = new ProductHierarchy(hiers);

            System.err.println("Creating cover...");
            // BaseHPM(hierarchy, prior, mincount, searchBreadth)

            TrainOptions options = new TrainOptions();
            DiscountOptions discountOptions = new DiscountOptions();

            BaseHPM oc = new BaseHPM(phier, pe, discountOptions);
            if (true) {
                System.err.println("Training cover...\n");

                long a = System.currentTimeMillis();
                oc.train(ds.getIDs(), pv, options);
                long b = System.currentTimeMillis();
                long t = b - a;

                System.err.println("Training time: " + t);
//
                oc.toDot2("/Users/mwh/Java/XHPM-MWH/output/pos.dot", false);
//

                System.err.println("CoverSize: " + oc.OCsize());
                System.err.println("BumpPost: " + oc.getBumpPostMDL(pv));


                OrderedCoverEstimator oce = new OrderedCoverEstimator(oc);

                for (int i : POSHier.getAllClasses()) {


                    HClass c1 = POSHier.getHClass(i);
                    if (!c1.equals(POSHier.getTop())) {
                        double sum = 0;
//                        System.err.println(c1.getName());
                        for (int j : POSHier.getAllClasses()) {
                            HClass c2 = POSHier.getHClass(j);
                            if (!c2.equals(POSHier.getTop())) {
                                System.err.print(c1.getName() + "\t" + c2.getName());
                                HClass[] classes = {c1, c2};
                                ProductHClass pc = new ProductHClass(phier, classes);
                                double cond = oce.getConditioningProb(pc);
                                double prob = oce.classAsDataPointEstimate(pc) - cond;
                                System.err.println("\t" + prob);
                                sum = sum + Math.pow(10, prob);
                            }
                        }
                        System.err.println("Sum: " + sum);
                    }
                }
//                NSuffixOCEstimator oce = new NSuffixOCEstimator(oc);

//                oce.save(args[2]);
//                PrintWriter out = new PrintWriter(new FileWriter(args[2] + ".log"));
//                out.println("Data points: " + oc.dataSize());
//                out.println("Train options: " + options);
//                out.println("");
//                out.println("Training time:" + t + "ms");
//                out.println("Cover size: " + oc.OCsize());
//                out.println("Score: " + oc.getBumpPostMDL(pv));
//                out.close();
            }

        }


        if (true) {

            TrainOptions options = new TrainOptions(args[0]);
            DiscountOptions discountOptions = new DiscountOptions();
            options.setIntOption("mincount", 2);
            options.setIntOption("depth", 10);
            options.setIntOption("rounds", 2);
            System.err.println(options);


            Storage storage = new MemoryStorage();
            String fn = args[1];
            DataSet ds = new DataSetStorage(storage);

            BufferedReader input = new BufferedReader(new FileReader(fn));

            // Just a counter to control how much data is readb
            int rl = 0;

            // List of words. Used to create empirical estimator
            List<String> wl = new ArrayList<String>();

            int maxread = Integer.MAX_VALUE;//Integer.parseInt(args[1]);

            int lc = 0;
            String line;
            System.err.println("Reading data...");

//            List<List<String>> ngrams = new ArrayList<List<String>>();
            while (((line = input.readLine()) != null) && rl < maxread) {
                rl++;
                line = "<s> <s> " + line + " </s>";
                String[] tokens = line.split(" ");

//                for (int i = 0; i < tokens.length - 1; i++) {
//                    List<String> ng = new ArrayList<String>(2);
//                    ng.add("^" + tokens[i]);
//                    ng.add("^" + tokens[i + 1]);
//                    ngrams.add(ng);
//                    ds.add(ng);
//                }

                for (int i = 0; i < tokens.length - 2; i++) {
                    List<String> ng = new ArrayList<String>(2);
                    ng.add("^" + tokens[i]);
                    ng.add("^" + tokens[i + 1]);
                    ng.add("^" + tokens[i + 2]);
//                    ngrams.add(ng);
                    ds.add(ng);
                }
                for (String t : tokens) {
                    // Add prefix to word
                    String w = "^" + t;
                    wl.add(w);
                }

                lc++;
            }
            input.close();

            System.err.println("Lines read: " + lc);
            System.err.println("words read: " + wl.size());
            System.err.println("");

            System.err.println("Creating prior...");
            EmpSuffix es = new EmpSuffix(wl);

            EmpSuffix ls = new LaPlaceSuffix(es);
            // Create estimator to handle expression
            EstimatorExp ees = new SimSuffixEstimatorExp(es, 100);
            EstimatorExp els = new SimSuffixEstimatorExp(ls, 100);



            // Array of estimators used in product-estimator
            //EstimatorExp[] ee = {se, se};
//            EstimatorExp[] ee = {els, ees};
            EstimatorExp[] ee = {els, els, ees};

            // Create product-estimator
            IndependentProductEst ie = new IndependentProductEst(ee);

            EstimatorExp pe = ie;

            System.err.println("Creating product dataset");
            // Make empty dataset

            System.err.println("DatasetSize: " + ds.size());

            // Create suffix-hiearchy
            SuffixHierarchy sufhier = new SuffixHierarchy();

            // Data for hierarchy is at index 0.

            SimpleView sufdv1 = new SimpleView(0, ds);
            SimpleView sufdv2 = new SimpleView(1, ds);
            SimpleView sufdv3 = new SimpleView(2, ds);

            // Create productView
//            View[] views = {sufdv1, sufdv2};
            View[] views = {sufdv1, sufdv2, sufdv3};
            ProductView pv = new ProductView(ds, views);

//            Hierarchy[] hiers = {sufhier, sufhier};
            Hierarchy[] hiers = {sufhier, sufhier, sufhier};

            // And now we make the suf x suf producthierarchy
            ProductHierarchy phier = new ProductHierarchy(hiers);


            // Create OC on the base of the hierarchy, the estimator and a minimun count
            System.err.println("Creating cover...");
            // BaseHPM(hierarchy, prior, mincount, searchBreadth)
            BaseHPM oc = new BaseHPM(phier, pe, discountOptions);
            if (true) {
                System.err.println("Training cover...\n");

                long a = System.currentTimeMillis();
                oc.train(ds.getIDs(), pv, options);
                long b = System.currentTimeMillis();
                long t = b - a;

                System.err.println("Training time: " + t);

                oc.toDot2(args[2] + ".dot", false);


                System.err.println("CoverSize: " + oc.OCsize());
                System.err.println("BumpPost: " + oc.getBumpPostMDL(pv));

                NSuffixOCEstimator oce = new NSuffixOCEstimator(oc);

                oce.save(args[2]);
                PrintWriter out = new PrintWriter(new FileWriter(args[2] + ".log"));
                out.println("Data points: " + oc.dataSize());
                out.println("Train options: " + options);
                out.println("");
                out.println("Training time:" + t + "ms");
                out.println("Cover size: " + oc.OCsize());
                out.println("Score: " + oc.getBumpPostMDL(pv));
                out.close();
            }
        }

    }
}